# menu_bar.py
import tkinter as tk
from tkinter import filedialog, messagebox

class MenuBar:
    def __init__(self, root, main_window=None, trainer=None):
        self.root = root
        self.main_window = main_window
        self.trainer = trainer  # 保存trainer引用
        self.create_menu()
        self.bind_shortcuts()

    def create_menu(self):
        # 创建菜单栏
        menubar = tk.Menu(self.root)
        self.root.config(menu=menubar)

        # 文件菜单
        file_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="文件", menu=file_menu)
        self.file_menu = file_menu
        file_menu.add_command(label="打开图片", command=self.open_image, accelerator="Ctrl+O")
        file_menu.add_command(label="打开视频", command=self.open_video, accelerator="Ctrl+V")
        file_menu.add_command(label="打开摄像", command=self.open_camera, accelerator="Ctrl+C")
        file_menu.add_command(label="保存结果", command=self.save_results, accelerator="Ctrl+S")
        file_menu.add_command(label="批量检测", command=self.batch_detect)
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.exit_app, accelerator="Ctrl+Q")

        # 训练菜单
        train_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="训练", menu=train_menu)
        train_menu.add_command(label="进入训练", command=self.train)
        train_menu.add_command(label="开始", command=self.start_training)
        train_menu.add_command(label="停止", command=self.stop_training)

        # 推理菜单
        infer_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="推理", menu=infer_menu)
        infer_menu.add_command(label="进入推理", command=self.predict)
        infer_menu.add_command(label="调整推理参数", command=self.adjust_infer_params)
        infer_menu.add_command(label="开始", command=self.start_inference, accelerator="F5")
        infer_menu.add_command(label="暂停", command=self.pause_inference, accelerator="F6")
        infer_menu.add_command(label="停止", command=self.stop_inference, accelerator="F7")

        # 关于菜单
        about_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="关于", menu=about_menu)
        about_menu.add_command(label="关于", command=self.show_about)

    def bind_shortcuts(self):
        # 绑定全局快捷键
        self.root.bind("<Control-o>", lambda e: self.open_image())
        self.root.bind("<Control-v>", lambda e: self.open_video())
        # self.root.bind("<Control-c>", lambda e: self.open_camera())
        self.root.bind("<Control-s>", lambda e: self.save_results())
        self.root.bind("<Control-q>", lambda e: self.exit_app())

        # F5-F7 快捷键
        self.root.bind("<F5>", lambda e: self.start_inference())
        self.root.bind("<F6>", lambda e: self.pause_inference())
        self.root.bind("<F7>", lambda e: self.stop_inference())

    # 菜单功能
    def open_image(self):
        if self.main_window:
            self.main_window.open_image()

    def open_video(self):
        if self.main_window:
            self.main_window.open_video()

    def open_camera(self):
        if self.main_window:
            self.main_window.open_camera()

    def save_results(self):
        if self.main_window:
            self.main_window.save_detection()

    def batch_detect(self):
        """批量检测功能"""
        try:
            # 打开文件选择对话框，允许选择多个文件
            file_paths = filedialog.askopenfilenames(
                title="选择图片文件进行批量检测",
                filetypes=[
                    ("图片文件", "*.jpg *.jpeg *.png *.bmp *.tiff *.JPG *.JPEG *.PNG *.BMP *.TIFF"),
                    ("所有文件", "*.*")
                ]
            )

            if not file_paths:
                return  # 用户取消选择

            # 将选中的文件路径传递给主窗口进行批量检测
            if self.main_window:
                # 清空当前路径输入框
                self.main_window.path_entry.delete(0, tk.END)
                # 在主窗口中设置批量文件路径（可以以分号分隔显示）
                self.main_window.path_entry.insert(0, ";".join(file_paths[:3]) + ("..." if len(file_paths) > 3 else ""))
            else:
                messagebox.showinfo("提示", f"已选择 {len(file_paths)} 个文件进行批量检测")

        except Exception as e:
            messagebox.showerror("错误", f"批量检测功能出错: {str(e)}")

    def exit_app(self):
        self.root.quit()

    def start_training(self):
        """开始训练"""
        if self.trainer:
            self.trainer.start_training()
        else:
            messagebox.showinfo("提示", "请先切换到训练界面")

    def stop_training(self):
        """停止训练"""
        if self.trainer:
            self.trainer.stop_training()
        else:
            messagebox.showinfo("提示", "当前没有正在进行的训练任务")

    def adjust_infer_params(self):
        if self.main_window:
            self.main_window.open_parameter_window()

    def start_inference(self):
        if self.main_window:
            self.main_window.start_detection()

    def pause_inference(self):
        if self.main_window:
            self.main_window.pause_detection()

    def stop_inference(self):
        if self.main_window:
            self.main_window.stop_detection()

    def show_about(self):
        messagebox.showinfo("关于", "这是一个基于 Ultralytics 的 AI 检测系统。\n版本：1.0\n作者：灵码")

    def predict(self):
        """进入推理界面"""
        try:
            # 关闭当前窗口
            self.root.destroy()

            # 导入并启动结果显示窗口
            from UI.ui.resultDisplay import main as start_result_display
            start_result_display()
        except Exception as e:
            messagebox.showerror("错误", f"无法启动推理界面: {str(e)}")

    def train(self):
        """进入训练界面"""
        try:
            # 关闭当前窗口
            self.root.destroy()

            # 导入并启动训练窗口
            from UI.ui.train import main as start_train_display
            start_train_display()
        except Exception as e:
            messagebox.showerror("错误", f"无法启动训练界面: {str(e)}")
